import os

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from matplotlib import rcParams
from matplotlib.ticker import MaxNLocator
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import cross_val_score


def ensure_picture_folder_exists():
    """Ensure the picture folder exists."""
    picture_folder = "./picture"
    if not os.path.exists(picture_folder):
        os.makedirs(picture_folder)
    return picture_folder


def set_matplotlib_params():
    """Set matplotlib parameters to avoid issues with fonts."""
    rcParams['font.sans-serif'] = ['Arial']
    rcParams['axes.unicode_minus'] = False


def load_data(filepath):
    """Load data from a CSV file."""
    try:
        data = pd.read_csv(filepath, encoding='ISO-8859-1')
        print(f"Data file '{filepath}' read successfully!")
        return data
    except FileNotFoundError:
        print(f"Error: The file '{filepath}' was not found. Please check the file path!")
        exit()
    except UnicodeDecodeError:
        print(
            f"Error: The file '{filepath}' encoding is not supported by ISO-8859-1. Please try another encoding format!")
        exit()


def preprocess_data(data):
    """Preprocess the data by filling missing values and converting categorical features."""
    # Fill missing values with 0
    data.fillna(0, inplace=True)
    # Convert country code (NOC) to numerical feature
    data['NOC'] = data['NOC'].astype('category').cat.codes
    return data


def extract_features(data):
    """Extract features and target variables from the data."""
    features = ['Year', 'NOC', 'Gold', 'Silver', 'Bronze', 'Total']
    X = data[features]
    y_gold = data['Gold']  # Target: Predict gold medals
    y_total = data['Total']  # Target: Predict total medals
    return X, y_gold, y_total


def visualize_data(data, picture_folder):
    """Visualize the data using various plots."""
    # 4.1 Distribution of medal counts
    plt.figure(figsize=(12, 6))
    sns.histplot(data['Gold'], bins=30, kde=True, color='blue', label='Gold Medals')
    sns.histplot(data['Total'], bins=30, kde=True, color='red', label='Total Medals')
    plt.title('Distribution of Gold and Total Medals')
    plt.xlabel('Medal Count')
    plt.ylabel('Frequency')
    plt.legend()
    plt.tight_layout()
    plt.grid(True)
    plt.savefig(f"{picture_folder}/1.3.1.png")
    plt.show()


def train_models(X, y_gold, y_total):
    """Train Random Forest Regressor models and evaluate their performance."""
    model_gold = RandomForestRegressor(n_estimators=100, random_state=42)
    model_total = RandomForestRegressor(n_estimators=100, random_state=42)

    # Cross-validation to evaluate model performance
    cv_scores_gold = cross_val_score(model_gold, X, y_gold, cv=5, scoring='neg_mean_squared_error')
    cv_scores_total = cross_val_score(model_total, X, y_total, cv=5, scoring='neg_mean_squared_error')

    print(f"Cross-validation MSE for gold medal prediction: {-cv_scores_gold.mean():.2f}")
    print(f"Cross-validation MSE for total medal prediction: {-cv_scores_total.mean():.2f}")

    # Train models
    model_gold.fit(X, y_gold)
    model_total.fit(X, y_total)

    return model_gold, model_total


def visualize_model_performance(model_gold, X, picture_folder):
    """Visualize the model performance and feature importance."""
    # 6.1 Feature importance for gold medal prediction
    plt.figure(figsize=(10, 6))
    plt.barh(X.columns, model_gold.feature_importances_, color='skyblue')
    plt.title('Feature Importance for Gold Medal Prediction')
    plt.xlabel('Importance')
    plt.ylabel('Feature')
    plt.tight_layout()
    plt.grid(axis='x')  # Only vertical grid lines
    plt.savefig(f"{picture_folder}/1.3.2.png")
    plt.show()


def predict_medals_for_2028(model_gold, model_total, data):
    """Predict the number of medals for the 2028 Olympics."""
    countries = data['NOC'].unique()
    predictions = []

    for country in countries:
        # Construct features for 2028
        X_2028 = pd.DataFrame({
            'Year': [2028],
            'NOC': [country],
            'Gold': [data[data['NOC'] == country]['Gold'].mean()],
            'Silver': [data[data['NOC'] == country]['Silver'].mean()],
            'Bronze': [data[data['NOC'] == country]['Bronze'].mean()],
            'Total': [data[data['NOC'] == country]['Total'].mean()]
        })

        # Predict gold medals and total medals
        gold_pred = model_gold.predict(X_2028)[0]
        total_pred = model_total.predict(X_2028)[0]

        predictions.append({
            'NOC': country,
            'Gold_2028': gold_pred,
            'Total_2028': total_pred
        })

    predictions_df = pd.DataFrame(predictions)
    return predictions_df


def generate_medal_table(predictions_df):
    """Sort predictions by gold medals and total medals to generate a medal table."""
    predictions_df = predictions_df.sort_values(by=['Gold_2028', 'Total_2028'], ascending=[False, False])
    print("Predicted medal table for the 2028 Los Angeles Olympics:")
    print(predictions_df)
    return predictions_df


def analyze_country_performance(predictions_df, data):
    """Analyze country performance by comparing historical medals with predicted medals."""
    historical_medals = data.groupby('NOC')[['Gold', 'Total']].mean().reset_index()
    historical_medals.rename(columns={'Gold': 'Historical_Gold', 'Total': 'Historical_Total'}, inplace=True)

    # Merge historical and predicted medals
    performance_df = pd.merge(predictions_df, historical_medals, on='NOC')

    # Calculate changes in medal counts
    performance_df['Gold_Change'] = performance_df['Gold_2028'] - performance_df['Historical_Gold']
    performance_df['Total_Change'] = performance_df['Total_2028'] - performance_df['Historical_Total']

    # Analyze improved and declined countries
    improved_countries = performance_df[performance_df['Gold_Change'] > 0]
    declined_countries = performance_df[performance_df['Gold_Change'] < 0]

    # Analyze countries winning medals for the first time
    first_time_winners = performance_df[
        (performance_df['Historical_Gold'] == 0) & (performance_df['Gold_2028'] > 0)]

    # Print analysis results
    print("Countries likely to improve:")
    print(improved_countries[['NOC', 'Gold_2028', 'Historical_Gold', 'Gold_Change']])

    print("Countries likely to decline:")
    print(declined_countries[['NOC', 'Gold_2028', 'Historical_Gold', 'Gold_Change']])

    print("Countries likely to win medals for the first time:")
    print(first_time_winners[['NOC', 'Gold_2028', 'Historical_Gold']])

    return performance_df


def visualize_predictions(predictions_df, performance_df, picture_folder):
    """Visualize the predictions and performance analysis."""
    # 10.1 Medal Table Visualization
    plt.figure(figsize=(12, 6))
    sns.barplot(x='Gold_2028', y='NOC', data=predictions_df.head(10), palette='viridis')
    plt.title('Predicted Medal Table for the 2028 Los Angeles Olympics (Top 10)')
    plt.xlabel('Gold Medals')
    plt.ylabel('Country')
    plt.tight_layout()
    plt.grid(axis='y')  # Only horizontal grid lines
    plt.savefig(f"{picture_folder}/1.3.3.png")
    plt.show()

    # 10.2 Country Performance Analysis Visualization
    plt.figure(figsize=(12, 6))
    performance_df['Gold_2028'] = performance_df['Gold_2028'].astype(int)

    sns.scatterplot(x='Historical_Gold', y='Gold_2028', hue='Gold_Change', data=performance_df, palette='coolwarm',
                    size='Gold_Change', sizes=(20, 200))
    plt.title('Comparison of Historical and Predicted Gold Medals')
    plt.xlabel('Historical Gold Medals')
    plt.ylabel('Predicted Gold Medals')

    # Force x-axis to display integer values
    plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))

    plt.axline((0, 0), slope=1, color='red', linestyle='--')  # Add diagonal line
    plt.tight_layout()
    plt.grid(True)
    plt.savefig(f"{picture_folder}/1.3.4.png")
    plt.show()


# Main execution
if __name__ == "__main__":
    # Ensure the picture folder exists
    picture_folder = ensure_picture_folder_exists()

    # Set matplotlib parameters
    set_matplotlib_params()

    # Step 1: Load data
    data = load_data(r'summerOly_medal_counts.csv')

    # Step 2: Data preprocessing
    data = preprocess_data(data)

    # Step 3: Feature extraction
    X, y_gold, y_total = extract_features(data)

    # Step 4: Data visualization
    visualize_data(data, picture_folder)

    # Step 5: Model training
    model_gold, model_total = train_models(X, y_gold, y_total)

    # Step 6: Model visualization
    visualize_model_performance(model_gold, X, picture_folder)

    # Step 7: Predict medals for 2028
    predictions_df = predict_medals_for_2028(model_gold, model_total, data)

    # Step 8: Generate medal table
    predictions_df = generate_medal_table(predictions_df)

    # Step 9: Analyze country performance
    performance_df = analyze_country_performance(predictions_df, data)

    # Step 10: Visualization
    visualize_predictions(predictions_df, performance_df, picture_folder)